package top.wshape1.shiziapi.message.websocket;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;
import top.wshape1.shiziapi.common.constant.RedisConstant;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;

/**
 * @author Wshape1
 * @version 1.0
 * @description
 * @since 2024-01-24
 */

@Slf4j
@Component
public class WsSessionManager {

    private static final ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 添加 session
     *
     * @param id
     */
    public void add(String id, WebSocketSession session) {
        SESSION_POOL.put(id, session);
        redisTemplate.opsForSet().add(RedisConstant.USER_ONLINE_KEY, id);
    }

    /**
     * 删除 session,会返回删除的 session
     *
     * @param id
     * @return
     */
    public WebSocketSession remove(String id) {
        WebSocketSession session = SESSION_POOL.remove(id);
        if (session != null)
            redisTemplate.opsForSet().remove(RedisConstant.USER_ONLINE_KEY, id);
        return session;
    }

    /**
     * 删除并同步关闭连接
     *
     * @param id
     */
    public void removeAndClose(String id) {
        WebSocketSession session = remove(id);
        if (session != null) {
            try {
                // 关闭连接
                if (session.isOpen())
                    session.close();
            } catch (IOException e) {
//                e.printStackTrace();
                log.error("关闭连接失败", e);
            }
        }
    }

    /**
     * 获得 session
     *
     * @param id
     * @return
     */
    public WebSocketSession get(String id) {
        return SESSION_POOL.get(id);
    }

    /**
     * 判断 session 是否存在
     *
     * @param id
     * @return
     */
    public boolean exist(String id) {
        return SESSION_POOL.containsKey(id);
    }

    /**
     * 查询redis判断该用户是否在线
     *
     * @param id
     * @return
     */
    public boolean isOnline(String id) {
        return Boolean.TRUE.equals(redisTemplate.opsForSet().isMember(RedisConstant.USER_ONLINE_KEY, id));
    }

    public void foreach(BiConsumer<String, WebSocketSession> consumer) {
        SESSION_POOL.forEach(consumer);
    }

}
